# K = 4, d = 4
# m_1 = 8, m_2 = m_3 = 20, m_4 = 18, n = 4
# C^(A) = 4, C^(B) = 2

# d大于等于4时，K > 1就不对，否则，无论K多大，都能相等
import numpy as np

K = 2
d = 4

lst_m = [4,4,4,4]
n = 2
C_A = 4
C_B = 2
lst_A = []
lst_B = []

for k in range(K):
	for i in range(d):
		A = np.random.normal(0, 1/np.sqrt(2*np.pi), (lst_m[i], C_A))
		B = np.random.normal(0, 1/np.sqrt(2*np.pi), (n, C_B))
		lst_A.append(A)
		lst_B.append(B)

X = np.random.uniform(0.0, 1.0, lst_m)


def normal_algo_1(I):
	lst_AB = []
	for k in range(K):
		for i in range(d):
			AB = np.kron(lst_A[k * d + i], lst_B[k * d + i])
			lst_AB.append(AB)

	lst_W = []
	for i in range(d):
		l_concat = []
		for k in range(K):
			l_concat.append(lst_AB[k * d + i])
		W = np.concatenate(l_concat, axis = -1)
		lst_W.append(W)
	C_W = C_A * C_B * K

	# reshape input to (m_{2}*m_{3}*...*m_{d}, m_{1})
	I = np.reshape(I, (lst_m[0], -1)).T
	for i in range(d):
		# ith CP core to (m_{i}, r^{A}*r^{B}, n_{i})
		W = np.transpose(np.reshape(lst_W[i], (lst_m[i], n, C_W)), (0, 2, 1))
		if i % 2 == 0:
			# O is (n_{1}*...*n_{i-1}m_{i+1}*...*m_{d}, r^{A}*r^{B}, n_{i})
			O = np.einsum('bm,mrn->brn', I, W)
			# reshape O to (n_{1}*...*n_{i-1}*n_{i}m_{i+2}*...*m_{d}, m_{i+1}, r^{A}*r^{B})
			if i == d - 1:
				O = np.reshape(O, (np.power(n, i), 1, -1, C_W, n))
			else:
				O = np.reshape(O, (np.power(n, i), lst_m[i + 1], -1, C_W, n))
			O = np.transpose(O, (0, 4, 2, 1, 3))
			if i == d - 1:
				O = np.reshape(O, (-1, 1, C_W))
			else:
				O = np.reshape(O, (-1, lst_m[i + 1], C_W))
				I = O
		if i % 2 != 0:
			# O is (n_{1}*...*n_{i-1}m_{i+1}*...*m_{d}, n_{i})
			O = np.einsum('bmr,mrn->bn', I, W)
			# reshape O to (n_{1}*...*n_{i-1}*n_{i}m_{i+2}*...*m_{d}, m_{i+1})
			if i == d - 1:
				O = np.reshape(O, (np.power(n, i), 1, -1, n))
			else:
				O = np.reshape(O, (np.power(n, i), lst_m[i + 1], -1, n))
			O = np.transpose(O, (0, 3, 2, 1))
			if i == d - 1:
				O = np.reshape(O, (-1, 1))
			else:
				O = np.reshape(O, (-1, lst_m[i + 1]))
				I = O
	Y = np.reshape(O, (-1))
	return Y


def single_algo_1(I, lst_W, C_W):
	# reshape input to (m_{2}*m_{3}*...*m_{d}, m_{1})
	I = np.reshape(I, (lst_m[0], -1)).T
	for i in range(d):
		# ith CP core to (m_{i}, r^{A}*r^{B}, n_{i})
		W = np.transpose(np.reshape(lst_W[i], (lst_m[i], n, C_W)), (0, 2, 1))
		if i % 2 == 0:
			# O is (n_{1}*...*n_{i-1}m_{i+1}*...*m_{d}, r^{A}*r^{B}, n_{i})
			O = np.einsum('bm,mrn->brn', I, W)
			# reshape O to (n_{1}*...*n_{i-1}*n_{i}m_{i+2}*...*m_{d}, m_{i+1}, r^{A}*r^{B})
			if i == d - 1:
				O = np.reshape(O, (np.power(n, i), 1, -1, C_W, n))
			else:
				O = np.reshape(O, (np.power(n, i), lst_m[i + 1], -1, C_W, n))
			O = np.transpose(O, (0, 4, 2, 1, 3))
			if i == d - 1:
				O = np.reshape(O, (-1, 1, C_W))
			else:
				O = np.reshape(O, (-1, lst_m[i + 1], C_W))
				I = O
		if i % 2 != 0:
			# O is (n_{1}*...*n_{i-1}m_{i+1}*...*m_{d}, n_{i})
			O = np.einsum('bmr,mrn->bn', I, W)
			# reshape O to (n_{1}*...*n_{i-1}*n_{i}m_{i+2}*...*m_{d}, m_{i+1})
			if i == d - 1:
				O = np.reshape(O, (np.power(n, i), 1, -1, n))
			else:
				O = np.reshape(O, (np.power(n, i), lst_m[i + 1], -1, n))
			O = np.transpose(O, (0, 3, 2, 1))
			if i == d - 1:
				O = np.reshape(O, (-1, 1))
			else:
				O = np.reshape(O, (-1, lst_m[i + 1]))
				I = O
	Y = np.reshape(O, (-1))
	return Y


def divided_algo_1(I):
	lst_AB = []
	for k in range(K):
		for i in range(d):
			AB = np.kron(lst_A[k * d + i], lst_B[k * d + i])
			lst_AB.append(AB)

	lst_Y = []
	for k in range(K):
		lst_W = []
		for i in range(d):
			lst_W.append(lst_AB[k * d + i])
		Y = single_algo_1(I, lst_W, C_A * C_B)
		lst_Y.append(Y)

	Y = lst_Y[0]
	for j in range(1, K):
		Y = Y + lst_Y[j]
	return Y


if __name__ == '__main__':
	Y_1 = normal_algo_1(X)
	print(Y_1.shape)
	print(Y_1)
	print('\n')
	Y_2 = divided_algo_1(X)
	print(Y_2.shape)
	print(Y_2)
